"""
调用dolphin转录音频为文本
"""

from threading import RLock
import dolphin
from typing import Union, Optional
from audios.dolphin_languages import get_language_code
from loguru import logger
from io import BytesIO
import re
from pathlib import Path
import torch
import tempfile
import warnings
warnings.filterwarnings('ignore')


class AudioAsr(object):

    lock = RLock()

    def __new__(cls, *args, **kwargs):
        with AudioAsr.lock:
            if not hasattr(AudioAsr, "_instance"):
                AudioAsr._instance = object.__new__(cls)
                cls.device = "cuda" if torch.cuda.is_available() else "cpu"
                cls.asr_model = dolphin.load_model('small', '/root/autodl-tmp/dolphin-small', cls.device)
                cls.reg = re.compile('<[\d\.]+>(.*)<[\d\.]+>')
        return AudioAsr._instance

    def recognize(self, audio_file: Union[Path, str],
                  zh_lang:str = '老挝语') -> Optional[str]:
        """
        调用dolphin模型转录音频为文字。
        """

        waveform = dolphin.load_audio(audio_file)
        lang_sym = get_language_code(zh_lang)
        try:
            res = self.asr_model(waveform, lang_sym)
        except Exception as e:
            logger.error(e)
            return None
        searched = self.reg.search(res.text)
        if searched is None:
            return None
        return searched.group(1)


if __name__ == '__main__':
    asr = AudioAsr()
    text = asr.recognize('../tests/audios/laos-djAjG3.mp3', zh_lang='老挝语')
    print(text)


